import yaml
from models.vit_moe import ViTMoE
from models.vit_shareparam import ViTMoEShareParam
import torch.nn as nn
from timm.models.vision_transformer import VisionTransformer

model_dict = {
    "ViTMoE": ViTMoE,
    "ViTMoEShareParam": ViTMoEShareParam,
    "VisionTransformer": VisionTransformer,
}


class args:
    # data
    data_root = "/home/ubuntu/ssk/mixture-of-experts/datasets/tiny-imagenet"
    dataset_name = "tiny-imagenet"
    batch_size = 256
    # Model 1 ViTMoE
    # moe_model = ViTMoE
    # model_kwargs = dict(
    #     embed_dim=256, depth=6, num_experts=4, num_classes=200, num_heads=8
    # )

    # Model 2 ViTMoEShareParam
    moe_model = ViTMoEShareParam
    model_kwargs = dict(
        embed_dim=256,
        depth=6,
        num_experts=7,
        num_classes=200,
        num_heads=8,
        share_dim=128,
    )

    # Model 3 ViT
    # moe_model = timm.models.VisionTransformer
    # model_kwargs = dict(embed_dim=256, depth=6, num_classes=200, num_heads=8)

    # loss func
    loss_func = nn.CrossEntropyLoss
    # optimizer
    lr = 4e-4
    expert_loss_coeff = 1
    weight_decay = 0.05
    # other
    epoch = 1500
    device = "cuda:1"

    def to_dict(self):
        return {k: str(v) for k, v in self.__dict__.items() if not k.startswith("__")}

    def from_yaml_file(self, file_path: str):
        with open(file_path, "r") as f:
            arguments: dict = yaml.load(f, Loader=yaml.FullLoader)
        for key in arguments.keys():
            setattr(self, key, arguments[key])

        self.moe_model = model_dict[self.moe_model]
        self.loss_func = getattr(nn, self.loss_func)

    def __repr__(self) -> str:
        return self.to_dict()
